from config import *
import numpy as np
import torch
import cv2

def load_img(img1,img2):

    img1 = cv2.imread(img1)[:,:,0]
    img1 = img1[:,cut_padding:-cut_padding].astype('float32')/255.0
    img2 = cv2.imread(img2)[:,:,0]
    img2 = img2[:,cut_padding:-cut_padding].astype('float32')/255.0

    cut_array1 = torch.from_numpy(img1).unsqueeze(0)
    cut_array2 = torch.from_numpy(img2).unsqueeze(0)
    img = torch.cat((cut_array1,cut_array2),0).unsqueeze(0)
    return img

def de_diag(acc , each_angle = False):
    result = np.sum(acc - np.diag(np.diag(acc)), 1) / 10.0
    if not each_angle:
        result = np.mean(result)
    return result

def eval(model):

    flag = 'test'

    model.cuda()
    model.eval()

    num_angle = len(test_angle_list)
    acc_table = np.zeros([len(test_probe_condition_list),num_angle,num_angle])

    count = 0
    for con,probe_condition in enumerate(test_probe_condition_list):
        for a1, probe_angle in enumerate(test_angle_list):
            for a2, gallery_angle in enumerate(test_angle_list):
                probe_person_num = 0
                probe_person_acc = 0
                for probe_identity in test_list:
                    for sub_probe_condition in probe_condition:
                        # print("Count/Total",count,'/363','Processing Probe Person',probe_person_num)
                        #找到每个probe具体每一个人
                        feature_list = []
                        candidate_list = []
                        for gallery_condition in test_gallery_condition_list:
                            for gallery_identity in test_list:
                                name1 = datapath+probe_identity+'/'+sub_probe_condition+'/'+probe_identity+'-'+sub_probe_condition+'-'+probe_angle+'.png'
                                name2 = datapath+gallery_identity+'/'+gallery_condition+'/'+gallery_identity+'-'+gallery_condition+'-'+gallery_angle+'.png'
                                if os.path.exists(name1) and os.path.exists(name2):
                                    img = load_img(name1,name2).cuda()
                                    feature = model(img,flag)
                                    feature_list.append(feature.item())
                                    candidate_list.append(gallery_identity)

                        if len(feature_list)>0:
                            print("Count/Total", count, '/363', 'Processing Probe Person', probe_person_num,name1,name2)
                            # print(name1,name2)
                            max_candidate_index = feature_list.index(max(feature_list))
                            max_candidate = candidate_list[max_candidate_index]
                            probe_person_num +=1
                            if max_candidate == probe_identity:
                                probe_person_acc += 1
                count += 1
                if probe_person_num ==0:
                    print(probe_condition,probe_angle,gallery_angle)
                    continue
                acc = probe_person_acc / probe_person_num
                acc_table[con,a1,a2] = acc

    print('===Rank-%d (Include identical-view cases)===' % (1))
    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
        np.mean(acc_table[0, :, :]),
        np.mean(acc_table[1, :, :]),
        np.mean(acc_table[2, :, :])))

    print('===Rank-%d (Exclude identical-view cases)===' % (1))
    a = de_diag(acc_table[0, :, :])
    b = de_diag(acc_table[1, :, :])
    c = de_diag(acc_table[2, :, :])
    # print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
    #     de_diag(acc_table[0, :, :]),
    #     de_diag(acc_table[1, :, :]),
    #     de_diag(acc_table[2, :, :])))
    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
        a,
        b,
        c))

    np.set_printoptions(precision=2, floatmode='fixed')
    print('===Rank-%d of each angle (Exclude identical-view cases)===' % (1))
    print('NM:', de_diag(acc_table[0, :, :], True))
    print('BG:', de_diag(acc_table[1, :, :], True))
    print('CL:', de_diag(acc_table[2, :, :], True))

    eval_reault = (a + b + c) / 3
    result_path = './result.txt'
    with open(result_path,'a') as f:
        f.write('NM={}'.format(a) + 'BG={}'.format(b) + 'CL={}'.format(c))
    return eval_reault

if __name__ == '__main__':
    checkpoint = torch.load(checkpoint_dir)
    model = checkpoint['model']
    print('Loading checkpoint from ',checkpoint_dir)
    result = eval(model)
    print("result", result)